package com.dahe.ysxt.vector.audit.service.impl;

import com.dahe.ysxt.vector.audit.dao.AuditDao;
import com.dahe.ysxt.vector.audit.dao.LikeAuditDao;
import com.dahe.ysxt.vector.audit.entity.AuditResult;
import com.dahe.ysxt.vector.audit.entity.LikeAuditResult;
import com.dahe.ysxt.vector.audit.service.AuditService;
import com.dahe.ysxt.vector.audit.thread.AuditThread;
import com.qcloud.cos.utils.CollectionUtils;
import com.tencent.tcvectordb.client.VectorDBClient;
import com.tencent.tcvectordb.model.Collection;
import com.tencent.tcvectordb.model.Database;
import com.tencent.tcvectordb.model.DocField;
import com.tencent.tcvectordb.model.Document;
import com.tencent.tcvectordb.model.param.collection.*;
import com.tencent.tcvectordb.model.param.database.ConnectParam;
import com.tencent.tcvectordb.model.param.dml.Filter;
import com.tencent.tcvectordb.model.param.dml.HNSWSearchParams;
import com.tencent.tcvectordb.model.param.dml.InsertParam;
import com.tencent.tcvectordb.model.param.dml.SearchByEmbeddingItemsParam;
import com.tencent.tcvectordb.model.param.entity.AffectRes;
import com.tencent.tcvectordb.model.param.entity.SearchRes;
import com.tencent.tcvectordb.model.param.enums.EmbeddingModelEnum;
import com.tencent.tcvectordb.model.param.enums.ReadConsistencyEnum;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;

@Service
public class AuditServiceImpl implements AuditService {
    private static final String DBNAME = "audit";
    private static final String COLL_NAME = "audit_result";
    private static final String COLL_NAME_ALIAS = "audit_result_alias";
    @Autowired
    private AuditDao auditDao;
    @Autowired
    private LikeAuditDao likeAuditDao;
    private Logger log = LoggerFactory.getLogger(this.getClass());

    final static VectorDBClient client =  new VectorDBClient(initConnectParam(), ReadConsistencyEnum.EVENTUAL_CONSISTENCY);
    @Override
    public int batchAuditDB() {
        ConnectParam connectParam = initConnectParam();
        VectorDBClient client = new VectorDBClient(connectParam, ReadConsistencyEnum.EVENTUAL_CONSISTENCY);
        //总条数
        int total = 0;
        //数据库中id最小的为1，最大的为381250，总共354393条 id基本是连续的
        int maxBatchNum = 383;
//        maxBatchNum =2 ;
        for(int batchNum = 1; batchNum < maxBatchNum; batchNum++) {
            //从DB数据库中每次ID递增1000查记录
            List<AuditResult> auditResultList = auditDao.queryAuditResults((batchNum - 1) * 1000 + 1, batchNum * 1000);
            total += auditResultList.size();
            if(CollectionUtils.isNullOrEmpty(auditResultList)) {
                //本批次未查到
                continue;
            }
            batchAuditVector(client, auditResultList);
            Date currentDate = new Date();
            // 创建SimpleDateFormat对象并定义时间格式
            SimpleDateFormat dateFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
            // 格式化当前时间并输出
            System.out.println("当前时间: " + dateFormat.format(currentDate) + "， 执行完第" + batchNum + "批,时间：" );
        }
        return total;
    }
    @Override
    public void batchAuditVector(VectorDBClient client, List<AuditResult> auditResultList) {
        Database database = client.database(DBNAME);
        Collection collection = database.describeCollection(COLL_NAME);
        List<Document> documentList = new ArrayList<Document>();
        for(AuditResult auditResult : auditResultList) {
            Document document = Document.newBuilder()
                    .withId(auditResult.getId() + "")
                    .addDocField(new DocField("bsm", auditResult.getBsm()))
                    .addDocField(new DocField("siteClassify", auditResult.getSiteClassify()))
                    .addDocField(new DocField("name", auditResult.getName()))
                    .addDocField(new DocField("title", auditResult.getTitle()))
                    .addDocField(new DocField("url", auditResult.getUrl()))
                    .addDocField(new DocField("urlMd5", auditResult.getUrlMd5()))
                    .addDocField(new DocField("conText", auditResult.getConText()))
                    .addDocField(new DocField("conTextMd5", auditResult.getConTextMd5()))
                    .addDocField(new DocField("errorWords", auditResult.getErrorWords()))
                    .addDocField(new DocField("correctWords", auditResult.getCorrectWords()))
                    .addDocField(new DocField("result", auditResult.getResult()))
                    .addDocField(new DocField("segment", auditResult.getConText()))
                    .build();
            documentList.add(document);
        }
        InsertParam insertParam = InsertParam.newBuilder().addAllDocument(documentList).withBuildIndex(true).build();
        collection.upsert(insertParam);
    }
    @Override
    public void createCollection( ) {
        ConnectParam connectParam = initConnectParam();
        VectorDBClient client = new VectorDBClient(connectParam, ReadConsistencyEnum.EVENTUAL_CONSISTENCY);
        createDatabaseAndCollection(client);
    }

    @Override
    public LikeAuditResult queryByEmbedding(Long id, String errorWords, String correctWords, String conText) {
//        ConnectParam connectParam = initConnectParam();
//        VectorDBClient client = new VectorDBClient(connectParam, ReadConsistencyEnum.EVENTUAL_CONSISTENCY);
        return queryByEmbedding(client, id, errorWords, correctWords, conText);
    }

    @Override
    public void batchSetEmbedding() {
        if(AuditThread.queue != null) {
            return;
        }
        log.info("开始初始化队列");
        AuditThread.queue = new ArrayBlockingQueue<AuditResult>(354393);
        int maxBatchNum = 383;
        for(int batchNum = 1; batchNum < maxBatchNum; batchNum++) {
            //从DB数据库中每次ID递增1000查记录
            List<AuditResult> auditResultList = auditDao.queryAuditResults((batchNum - 1) * 1000 + 1, batchNum * 1000);
            if(CollectionUtils.isNullOrEmpty(auditResultList)) {
                //本批次未查到
                continue;
            }
            for(AuditResult auditResult : auditResultList) {
                AuditThread.queue.add(auditResult);
            }
            log.info("初始化完第{}批", batchNum);
        }
        /*启20个线程查询*/
        for(int i=0; i < 20; i++) {
            AuditThread auditThread = new AuditThread();
            auditThread.start();
            log.info("第{}个线程启动完毕", i + 1);
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }


    /**
     * 相似性检索
     * @param client
     */
    private static LikeAuditResult queryByEmbedding(VectorDBClient client, Long id, String errorWords, String correctWords, String conText) {
        Database database = client.database(DBNAME);
        Collection collection = database.describeCollection(COLL_NAME);
        SearchByEmbeddingItemsParam searchByEmbeddingItemsParam = SearchByEmbeddingItemsParam.newBuilder()
                .withEmbeddingItems(Arrays.asList(conText))
                // 若使用 HNSW 索引，则需要指定参数 ef，ef 越大，召回率越高，但也会影响检索速度
                .withParams(new HNSWSearchParams(200))
                // 设置标量字段的 Filter 表达式，过滤所需查询的文档
                .withRetrieveVector(false)
                // 指定 Top K 的 K 值
                .withLimit(2)
                // 使用 filter 过滤数据
                .withFilter(new Filter("  errorWords = \"" + errorWords + "\" and correctWords = \"" + correctWords + "\" "))
                // 指定返回的 fields
                .withOutputFields(Arrays.asList(  "conText", "conTextMd5", "result", "errorWords", "correctWords"))
                .build();
        SearchRes searchRes = collection.searchByEmbeddingItems(searchByEmbeddingItemsParam);
        List<List<Document>> siDocs = searchRes.getDocuments();
        LikeAuditResult likeAuditResult = new LikeAuditResult();
        List<Document> docs = siDocs.get(0);
        likeAuditResult.setId(id);
        likeAuditResult.setErrorWords(errorWords);
        likeAuditResult.setCorrectWords(correctWords);
        likeAuditResult.setConText(conText);
        for(Document document : docs) {
            //本身
            if(likeAuditResult.getId().toString().equals(document.getId())) {
                continue;
            }
            likeAuditResult.setLikeId(Long.parseLong(document.getId()));
            likeAuditResult.setLikeScore(document.getScore());
            for(DocField docField : document.getDocFields()) {
                if("result".equals(docField.getName())) {
                    likeAuditResult.setLikeAuditResult((String)docField.getValue());
                    continue;
                }
                if("conText".equals(docField.getName())) {
                    likeAuditResult.setLikeConText((String)docField.getValue());
                    continue;
                }
                if("conTextMd5".equals(docField.getName())) {
                    likeAuditResult.setLikeConTextMd5((String)docField.getValue());
                    continue;
                }

            }

        }

        return likeAuditResult;
    }

    private static ConnectParam initConnectParam() {
        System.out.println("\tvdb_url: " + System.getProperty("vdb_url"));
        System.out.println("\tvdb_key: " + System.getProperty("vdb_key"));
        /*低配版*/
//        System.setProperty("vdb_url", "http://lb-rmfsvvoh-84vj9k7fcdan3o2x.clb.ap-beijing.tencentclb.com:40000");
//        System.setProperty("vdb_key", "vepGVsPr9ddIZgmvhaw5T0C5NOrKdSaee2MH4UA1");
        /*高配版*/
        System.setProperty("vdb_url", "http://lb-lq5x9sm9-xl29rc643uyb5hgs.clb.ap-beijing.tencentclb.com:20000");
        System.setProperty("vdb_key", "UV0CykJDvvbE55p5jIu6i6teicnKrONxLmr07ePh");

        return ConnectParam.newBuilder()
                .withUrl(System.getProperty("vdb_url"))
                .withUsername("root")
                .withKey(System.getProperty("vdb_key"))
                .withTimeout(30)
                .build();
    }

    public  void createDatabaseAndCollection(VectorDBClient client) {
        // 1. 创建数据库
        Database db = client.createDatabase(DBNAME);

        // 3. 创建 collection
        CreateCollectionParam collectionParam = initCreateCollectionParam(COLL_NAME);
        db.createCollection(collectionParam);

        // 5. 设置 collection 别名
        System.out.println("---------------------- setAlias ----------------------");
        AffectRes affectRes = db.setAlias(COLL_NAME, COLL_NAME_ALIAS);
        System.out.println("\tres: " + affectRes.toString());

    }

    private static CreateCollectionParam initCreateCollectionParam(String collName) {
        // Embedding模型
        Embedding embedding = Embedding
                .newBuilder()
                .withModel(EmbeddingModelEnum.BGE_BASE_ZH)
                .withField("segment")
                .withVectorField("vector")
                .build();

        return CreateCollectionParam.newBuilder()
                .withName(collName)
                .withShardNum(3)
                .withReplicaNum(2)
                //免费版副本数必须是0
//                .withReplicaNum(0)
                .withDescription("test collection0")
                .addField(new FilterIndex("id", FieldType.String, IndexType.PRIMARY_KEY))
                .addField(new VectorIndex("vector", EmbeddingModelEnum.BGE_BASE_ZH.getDimension(), IndexType.HNSW,
                        MetricType.COSINE, new HNSWParams(16, 200)))
                .addField(new FilterIndex("segment", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("name", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("title", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("bsm", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("url", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("conTextMd5", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("conText", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("siteClassify", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("urlMd5", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("correctWords", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("errorWords", FieldType.String, IndexType.FILTER))
                .addField(new FilterIndex("result", FieldType.String, IndexType.FILTER))
                .withEmbedding(embedding)
                .build();
    }

}
